package at.ltasat.classifier.knn;

import java.io.IOException;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import at.ltasat.classifier.ObjectFeaturesFactory;
import at.ltasat.classifier.ClazzManager;
import at.ltasat.classifier.Classification;
import at.ltasat.classifier.Classifier;
import at.ltasat.classifier.Clazz;
import at.ltasat.classifier.Membership;

final class KNearestNeighbourClassifier<O, D> extends Classifier<O, D> {
	private final int k;
	private final int factor;
	private final Map<Point, Membership<D>[]> data;

	public KNearestNeighbourClassifier(String name, ClazzManager<D> classManager, ObjectFeaturesFactory<O> factory, int k, int factor, Map<Point, Membership<D>[]> data) {
		super(name, classManager, factory);
		assert k >= 1;
		assert factor >= 1;
		assert data != null;
		this.k = k;
		this.factor = factor;
		this.data = data;
	}

	public int getK() {
		return k;
	}

	public int getFactor() {
		return factor;
	}

	public int getDataSize() {
		return data.size();
	}

	public Set<Point> getDataPoints() {
		return data.keySet();
	}

	public Membership<?>[] getData(Point point) {
		return data.get(point).clone();
	}

	@Override
	public Classification<O, D> classify(final O object) {
		assert object != null;
		if (data.size() == 0) {
			return null;
		}
		Point point = Point.create(getObjectFeatures(object), factor);
		List<Point> neighbourhood = getNeighbourhood(point);
		Map<Clazz<D>, Double> memberships = getMemberships(point, neighbourhood);
		assert memberships.size() > 0;
		return new Classification<O, D>(object, memberships);
	}

	private List<Point> getNeighbourhood(final Point point) {
		List<Point> neighbourhood = new LinkedList<Point>();
		for (Point current : data.keySet()) {
			double distanceToCurrent = point.getDistanceTo(current);
			int index = 0;
			for (Point neighbour : neighbourhood) {
				double distanceToNeighbour = point.getDistanceTo(neighbour);
				if (distanceToCurrent < distanceToNeighbour) {
					index++;
				} else {
					break;
				}
			}
			neighbourhood.add(index, current);
			if (neighbourhood.size() > k) {
				neighbourhood.remove(0);
			}
		}
		return neighbourhood;
	}

	private Map<Clazz<D>, Double> getMemberships(Point point, List<Point> neighbourhood) {
		Map<Clazz<D>, Double> weightedMemberships = getWeightedMemberships(point, neighbourhood);
		return normalize(weightedMemberships);
	}

	private static double getDistanceSum(final Point point, final Iterable<Point> points) {
		double distance = 0;
		for (Point neighbour : points) {
			distance += point.getDistanceTo(neighbour);
		}
		return distance;
	}

	private static <D> Map<Clazz<D>, Double> normalize(final Map<Clazz<D>, Double> weightedMemberships) {
		final double factor = getNormalizationFactor(weightedMemberships);
		final Map<Clazz<D>, Double> normalizedMemberships = new HashMap<Clazz<D>, Double>();
		for (Clazz<D> clazz : weightedMemberships.keySet()) {
			normalizedMemberships.put(clazz, weightedMemberships.get(clazz) / factor);
		}
		return normalizedMemberships;
	}

	private static <D> double getNormalizationFactor(final Map<Clazz<D>, Double> weightedMemberships) {
		double factor = 0;
		for (Clazz<D> clazz : weightedMemberships.keySet()) {
			factor += weightedMemberships.get(clazz);
		}
		return factor;
	}

	@SuppressWarnings("unchecked")
	private Map<Clazz<D>, Double> getWeightedMemberships(final Point point, final Iterable<Point> neighbourhood) {
		final double distance = getDistanceSum(point, neighbourhood);
		final Map<Clazz<D>, Double> memberships = new HashMap<Clazz<D>, Double>();
		for (Point neighbour : neighbourhood) {
			final double myDistance = point.getDistanceTo(neighbour);
			final double weight = (distance == 0) ? (1) : (1 - (myDistance / distance));
			for (Membership<?> membership : data.get(neighbour)) {
				Clazz<D> clazz = ((Membership<D>)membership).getClazz();
				if (!memberships.containsKey(clazz)) {
					memberships.put(clazz, 0.0);
				}
				memberships.put(clazz, memberships.get(clazz) + membership.getDegree() * weight);
			}
		}
		return memberships;
	}

	@Override
	public void write(final OutputStream out) throws IOException {
		KNearestNeighbourClassifierFactory.writeClassifier(out, this);
	}

}
